#!/usr/bin/env python
# -*- encoding: UTF-8 -*-
from __future__ import print_function

import numpy as np
import mpl_toolkits.mplot3d.axes3d as p3
import matplotlib.pyplot as p   # for figure
import cv2
from robot import tr2rpy
from numpy.linalg import inv
import math
from Queue import deque
import pickle

from common import draw_str,thinning,get_line

class Cam(object):
    def __init__(self,name=None,index=0):
        self.unit = 20.0 # 标定板单位为 20mm
        self.height = 480 # 相机标定时的高
        self.width = 640 # 相机标定时的宽
        self.ratio = 1 # 图像缩放比
        self.tr = None # 从标定板坐标系到相机坐标系的变换矩阵
        self.tr2 = None # 从另一个相机变换到本相机的变换矩阵, self.tr2 = self.tr*inv(other.tr)
        self._tr2 = None # 用来保存相对矩阵的历史值,用来均值滤波

        self.name = name
        self.mtx = self.get_cameraMatrix()
        self.dist = self.get_distCoeffs()
        self.cap = cv2.VideoCapture(index)

    def __getattr__(self, item):
        if item in self.__dict__:
            return self.__dict__[item]
        if item == "P": # 相机矩阵, P_{3x4}=(K_{3,3},0)_{3,4}
            # 将本相机坐标系中的三维坐标变换到本相机的二维图像中
            K = self.get_cameraMatrix()*self.ratio
            C = np.matrix(np.zeros((3,1)))
            P = np.hstack((K,C))
            return P
        elif item == "Pw": # 相机矩阵, P_{3x4}=K_{3,3}*R_{3,3}*(I,-C)_{3,4}
            if self.tr is None: return None
            # 将标定板坐标系中的三维坐标变换到本相机的二维图像中
            K = self.get_cameraMatrix()*self.ratio
            R = self.tr[:3,:3]
            I = np.eye(3)
            C = self.tr[:3,3]
            P = K*R*np.hstack((I,C))
            return P
        elif item == "P2": # 相机矩阵, P_{3x4}=K_{3,3}*R_{3,3}*(I,-C)_{3,4}
            if self.tr2 is None: return None
            # 将另一个相机坐标系中的三维坐标变换到本相机的二维图像中
            K = self.get_cameraMatrix()*self.ratio
            R = self.tr2[:3,:3]
            I = np.eye(3)
            C = self.tr2[:3,3]
            P = K*R*np.hstack((I,C))
            return P

    def findTr(self,src=None,w=9,h=4,unit=15,reverse=False):
        '''
        从标定板获得从世界坐标系到相机坐标系的变换矩阵
        :param src:  标定图像
        :param w:  标定角点的宽点个数
        :param h:  标定角点的高点个数
        :param unit: 标定板的单位
        :param reverse:  是否翻转标定板的世界坐标
        :return: Tr 从世界坐标系到相机坐标系的表换矩阵
        '''
        if src is None: src = self.get_img()
        if len(src.shape) == 3: height,width,depth = src.shape
        elif len(src.shape) == 2: height,width = src.shape
        else: raise ValueError

        assert width/self.width == height/self.height
        self.ratio = width/float(self.width)
        objp = np.zeros((h*w,3), np.float32)
        objp[:,:2] = np.mgrid[0:w,0:h].T.reshape(-1,2)
        objp = objp*unit/self.unit
        if reverse:objp = objp[::-1]
        criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
        gray = cv2.cvtColor(src,cv2.COLOR_BGR2GRAY)
        found, corners = cv2.findChessboardCorners(gray, (w,h),None)
        if not found: return None
        cv2.cornerSubPix(gray,corners,(11,11),(-1,-1),criteria)
        cv2.drawChessboardCorners(src,(w,h),corners,found)
        mtx = self.get_cameraMatrix()*self.ratio # 相机参数矩阵要等比例缩放
        dist = self.get_distCoeffs() # 畸变矫正参数与图像缩放无关
        ret,rvec,tvec = cv2.solvePnP(objp, corners,mtx,dist)
        if not ret: return None
        tr = np.matrix(np.eye(4))
        tr[:3,:3] = cv2.Rodrigues(rvec)[0]
        tr[:3,3] = tvec * self.unit # 换算成 mm
        x,y,z = np.array(tr[:3,3]).flatten()
        wz,wy,wx = np.array(tr2rpy(tr)).flatten()
        draw_str(src,(10,10),"x=%f, y=%f, z=%f" %(x,y,z))
        draw_str(src,(10,25),"wx=%f, wy=%f, wz=%f" %(wx,wy,wz))
        cv2.imshow(self.name+"tr",src)
        self.tr = tr
        return tr

    def findTr2(self,cam2,src1=None,src2=None,w=9,h=4,unit=15):
        tr = self.findTr(src1,w=w,h=h,unit=unit)
        tr2 = cam2.findTr(src2,w=w,h=h,unit=unit)
        if tr is None or tr2 is None: return None
        tr2 = tr*inv(tr2)
        if self._tr2_list is None:
            self.tr2 = tr2
            self._tr2 = deque()
            self._tr2.append(tr2)
        elif len(self._tr2_list) < 10:
            k = len(self._tr2)
            self.tr2 = (k*self.tr2+tr2)/(k+1.0)
            self._tr2.append(tr2)
        else:
            old = self._tr2.popleft()
            self.tr2 = (10*self.tr2-old + tr2)/10.0
            self._tr2.append(tr2)
        return self.tr2

    def get_cameraMatrix(self):
        mtx = np.eye(3)

        if self.name == '+':
            fx = 1511.64442682
            fy = 1514.72927917
            cx = 331.982535978
            cy = 268.613047699
        elif self.name == "++":
            fx = 1511.64442682
            fy = 1514.72927917
            cx = 331.982535978
            cy = 268.613047699
        else:
            fx = 1500.0
            fy = 1500.0
            cx = self.width/2.0
            cy = self.height/2.0

        mtx[0,0],mtx[1,1],mtx[0,2],mtx[1,2] = fx,fy,cx,cy
        return mtx

    def get_distCoeffs(self):
        dist = np.zeros((1,5))

        if self.name == "+":
            k1 = -0.105356293294
            k2 = 3.25067090251
            p1 = 0.00155905212496
            p2 = -0.00569751368371
            k3 = -25.6347944753
        elif self.name == "++":
            k1 = -0.105356293294
            k2 = 3.25067090251
            p1 = 0.00155905212496
            p2 = -0.00569751368371
            k3 = -25.6347944753
        else:
            k1 = 0
            k2 = 0
            p1 = 0
            p2 = 0
            k3 = 0

        dist[0] = k1,k2,p1,p2,k3
        return dist

    def get_img(self):
        ret, frame = self.cap.read()
        if ret:
            self._frame = frame
        elif hasattr(self,"_frame"):
            frame = self._frame
        else:
            frame = self.get_img()
        cv2.imshow(self.name,frame)
        return frame.copy()

    def get_line(self):
        src = self.get_img()
        return get_line(src,self.mtx,self.dist,win_name=self.name)

    def get_flat(self):
        '''
        返回在图像中找到的直线,反投影在世界坐标系下的平面,
        的齐次表示
        '''
        line = self.get_line()
        if line is None or self.tr is None: return None
        flat_cam = self.P.transpose()*line
        tr = inv(self.tr) # 将相机坐标系变换到世界坐标系的矩阵
        flat = tr*flat_cam # 此时计算出来的平面最后一位为 0 表示过原点, 不是正确的结果
        C = tr[:,3] #  相机坐标系原点在世界坐标系下的表示
        flat[-1] = -C.transpose()*flat # 确保相机原点在该平面上
        return flat

def dual_plucker(L):
    """
    一般 plucker 表示是用两空间点的齐次坐标表示, L_{4,4} = AB^T - BA^T
    但也可以用两空间的平面的齐次坐标表示, L^* = pi1*pi2^T - pi2*pi1^T
    两者互为对偶关系: L^* x L = 0, 关系如下:
    l_12 : l_13 : l_14 : l_23 : l_42 : l_34 ==
    == l^*_34 : l^*_42 : l^*_23 : l^*_14 : l^*_13 : l^*_12
    :param L: 空间直线的 plucker 表示矩阵
    :return: 相同空间直线的 plucker 表示矩阵的对偶表示矩阵
    """
    dual_L = L.copy()
    dual_L[0,1] = L[2,3]
    dual_L[0,2] = L[3,1]
    dual_L[0,3] = L[1,2]
    dual_L[1,2] = L[0,3]
    dual_L[3,1] = L[0,2]
    dual_L[2,3] = L[0,1]
    return dual_L - dual_L.transpose()

def plucker_line(cam1,cam2):
    """
    空间直线在 标定板 坐标系下的 plucker 矩阵表示: L_{4,4} = AB^T - BA^T
    A,B 为在 cam1 下直线 L 通过的两个空间点的齐次表示
    :param cam1: 相机1
    :param src1: 相机1的图像
    :param cam2: 相机2
    :param src2: 相机2的图像
    :return: L, Lstar  在标定板坐标系下的 pluker 矩阵: 两点式矩阵 L 和两平面式矩阵 Lstar
    """
    flat1 = cam1.get_flat()
    flat2 = cam2.get_flat()
    if flat1 is None or flat2 is None: return None,None
    Lstar = flat1*flat2.transpose() - flat2*flat1.transpose()
    L = dual_plucker(Lstar)
    return L, Lstar

class twoCam:
    def __init__(self):
        self.cam1 = Cam("+",0)
        self.cam2 = Cam("++",1)

    def __getattr__(self, item):
        if item == "Tr":
            return self.cam1.findTr2(self.cam2)

    def plucker_line(self):
        return plucker_line(self.cam1,self.cam2)

def main():
    two_cam = twoCam()
    while True:
        ch = cv2.waitKey(1) & 0xFF
        if ch == 27:
            break
        elif ch == ord('s'):
            cv2.imwrite(two_cam.cam1.name+".jpg",two_cam.cam1.get_img())
            cv2.imwrite(two_cam.cam2.name+".jpg",two_cam.cam2.get_img())
            pickle.dump(two_cam.cam1.tr,open("tr1.pkl", 'wb'))
            pickle.dump(two_cam.cam2.tr,open("tr2.pkl", 'wb'))

        Tr = two_cam.Tr
        L,Lstar = two_cam.plucker_line()
        if not (L is None):
            print("L:",L)
            print("Lstar:",Lstar)
        
        if Tr is None : continue
        x,y,z = np.array(Tr[:3,3]).flatten()
        wz,wy,wx = np.array(tr2rpy(Tr)).flatten()
        print("x=%f, y=%f, z=%f, wx=%f, wy=%f, wz=%f" %(x,y,z,wx,wy,wz))

if __name__ == '__main__':
    main()